# encoding:utf8
from datetime import datetime
import backtrader as bt
import numpy as np
import pandas as pd
from loguru import logger
from scipy.stats import linregress


class SmaCross(bt.Strategy):
    # list of parameters which are configurable for the strategy
    params = dict(
        pfast=10,  # period for the fast moving average
        pslow=30  # period for the slow moving average
    )

    def __init__(self):
        sma1 = bt.ind.SMA(period=self.p.pfast)  # fast moving average
        print(sma1)
        sma2 = bt.ind.SMA(period=self.p.pslow)  # slow moving average
        self.crossover = bt.ind.CrossOver(sma1, sma2)  # crossover signal
        print(self.crossover)

        self.sma1 = sma1
        # self.signal_add(bt.SIGNAL_LONG, self.crossover)  # use it as LONG signal
        self.dataclose = self.datas[0].close
        print(self.datas)
        logger.info(self.dataclose)

    def next(self):
        print(self.crossover[0])
        print(self.datas[0], self.sma1[0])
        if not self.position:  # not in the market
            if self.crossover > 0:  # if fast crosses slow to the upside
                self.buy()  # enter long

        elif self.crossover < 0:  # in the market & cross to the downside
            self.close()  # close long position

    def log(self, txt, dt=None):
        ''' Logging function for this strategy'''
        dt = dt or self.datas[0].datetime.date(0)
        # print('%s, %s' % (dt.isoformat(), txt))


def momentum_func(self, price_array):
    r = np.log(price_array)
    slope, _, rvalue, _, _ = linregress(np.arange(len(r)), r)
    annualized = (1 + slope) ** 252
    return (annualized * (rvalue ** 2))


# 定义动量指标
class Momentum(bt.ind.OperationN):
    lines = ('trend',)
    params = dict(period=90)
    func = momentum_func


class DummyInd(bt.Indicator):
    lines = ('dummyline',)

    params = (('value', 5),)

    def __init__(self):
        self.lines.dummyline = bt.Max(0.0, self.params.value)


# https://zhuanlan.zhihu.com/p/321149887?ivk_sa=1024320u
class BacktraderEngine:
    def __init__(self, init_cash=100000.0):
        self.init_cash = init_cash
        self.feed = CSVDatafeed()
        cerebro = bt.Cerebro()
        cerebro.broker.setcash(init_cash)

        self.cerebro = cerebro

    def add_data(self, code, start=datetime(2010, 1, 1), end=datetime.now().date()):
        # 加载数据
        self.feed.add_data(code, DATA_DIR_CSV.joinpath('{}.csv'.format(code)))
        df = self.feed.get_df(code)
        df = to_backtrader_dataframe(df)
        data = bt.feeds.PandasData(dataname=df, name=code, fromdate=start, todate=end)

        self.cerebro.adddata(data)  # Add the data feed

        # self.cerebro.addanalyzer(bt.analyzers.TimeReturn, timeframe=bt.TimeFrame.Years,
        #                         data=data, _name='mydata_return')
        self.cerebro.addobserver(bt.observers.Benchmark,
                                 data=data)

    def add_stragegy(self, strategy):
        self.cerebro.addstrategy(strategy)

    def run(self):
        # self.cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='_TimeReturn')
        self.cerebro.addanalyzer(bt.analyzers.Returns, _name='收益')
        self.cerebro.addanalyzer(bt.analyzers.TradeAnalyzer)
        self.cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='AnnualReturn')
        self.cerebro.addanalyzer(bt.analyzers.SharpeRatio, riskfreerate=0.0, annualize=True, _name='SharpeRatio')
        self.cerebro.addanalyzer(bt.analyzers.DrawDown, _name='DrawDown')

        self.cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')

        self.results = self.cerebro.run()  # run it all

    def analysis(self):
        # 获取回测结束后的总资金
        portfolio_stats = self.results[0].analyzers.getbyname('pyfolio')
        returns, positions, transactions, gross_lev = portfolio_stats.get_pf_items()
        returns.index = returns.index.tz_convert(None)

        import quantstats
        df = self.feed.get_df('000300.SH')
        df= df[['rate']]
        df.index = pd.to_datetime(df.index)
        print(df)

        quantstats.reports.html(returns, benchmark=df, output='stats.html', title='Stock Sentiment')



        portvalue = self.cerebro.broker.getvalue()
        pnl = portvalue - self.init_cash
        # 打印结果
        print(f'总资金: {round(portvalue, 2)}')
        print('收益', pnl)
        accu_return = portvalue / self.init_cash - 1
        logger.info('总收益率:{}'.format(accu_return))

        annu_ret = round((accu_return + 1) ** (252 / (12 * 252 + 252 * 3 / 4)) - 1, 3)
        logger.info('年化收益:{}'.format(annu_ret))

        strats = [x for x in self.results]  # 取得两个策略的运行结果

        for ret in strats:
            print("--------------- AnnualReturn -----------------")
            print(ret.analyzers.AnnualReturn.get_analysis())
            print("--------------- SharpeRatio -----------------")
            print(ret.analyzers.SharpeRatio.get_analysis())
            print("--------------- DrawDown -----------------")
            print(ret.analyzers.DrawDown.get_analysis())

            # tdata_analyzer = ret.analyzers.getbyname('datareturns')
            # print(tdata_analyzer.get_analysis())

        # self.cerebro.plot()  # and plot it with a single command
        # self.cerebro.plot(volume=False)

        from backtrader_plotting import Bokeh
        from backtrader_plotting.schemes import Tradimo

        plotconfig = {
            'id:ind#0': dict(
                subplot=True,
            ),
        }
        b = Bokeh(style='line', scheme=Tradimo(), plotconfig=plotconfig)
        self.cerebro.plot(b)
        # self.cerebro.plot(volume=False)


from engine.config import DATA_DIR_CSV
from engine.data_utils import to_backtrader_dataframe

from engine.stragegy.strategy_rotation import StrategyRotation
from engine.stragegy.stragegy_buyhold import StratgeyBuyHold


# 策略选择类
class StFetcher(object):
    _STRATS = [StratgeyBuyHold, StrategyRotation]  # 注册策略

    def __new__(cls, *args, **kwargs):
        idx = kwargs.pop('idx')  # 策略索引

        obj = cls._STRATS[idx](*args, **kwargs)
        return obj

if __name__ == '__main__':
    from engine.datafeed.datafeed_csv import CSVDatafeed
    e =BacktraderEngine()


    e.add_data('000300.SH')
    #e.add_data('000905.SH')
    #e.add_data('399006.SZ')
    e.add_data('SPX')
    #e.add_data('N225')
    #e.add_stragegy(StratgeyBuyHold)

    #e.cerebro.optstrategy(StFetcher, idx=[0, 1])

    e.add_stragegy(StrategyRotation)
    e.run()
    e.analysis()
